# -*- coding: utf-8 -*-
"""
NIPT的时点选择与胎儿的异常判定 - 问题一分析代码
作者: 数学建模团队
日期: 2025年9月
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import statsmodels.api as sm
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import r2_score
from pygam import LinearGAM, s, te
from scipy import stats
from statsmodels.stats.outliers_influence import variance_inflation_factor
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持和LaTeX渲染
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']  # 使用黑体和微软雅黑
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['mathtext.fontset'] = 'stix'  # 使用STIX字体，支持数学符号

# 1. 数据读取与预处理
print("="*60)
print("1. 数据读取与预处理")
print("="*60)

# 读取数据
df = pd.read_excel('处理后数据.xlsx')

# 数据预处理 - 选择需要的列并重命名
data = df[['检测孕周(处理后)', '孕妇BMI（处理后）', 'Y染色体浓度（处理后）']].copy()
data.columns = ['week', 'bmi', 'y_concentration']  # 重命名为英文便于编码

# 删除缺失值
data = data.dropna()
print(f"数据集大小: {data.shape}")
print(f"孕周范围: {data['week'].min():.2f} - {data['week'].max():.2f}")
print(f"BMI范围: {data['bmi'].min():.2f} - {data['bmi'].max():.2f}")
print(f"Y染色体浓度范围: {data['y_concentration'].min():.4f} - {data['y_concentration'].max():.4f}")

# 2. 探索性数据分析
print("\n" + "="*60)
print("2. 探索性数据分析")
print("="*60)

# 计算相关系数矩阵
corr_matrix = data.corr()
print("相关系数矩阵:")
print(corr_matrix)

# 绘制热力图
plt.figure(figsize=(10, 8))
sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0, square=True, fmt='.3f')
plt.title('变量间相关性热力图', fontsize=16)
plt.tight_layout()
plt.savefig('correlation_heatmap.png', dpi=300, bbox_inches='tight')
plt.show()

# 绘制散点图（按BMI分组）
# 定义BMI分组
bins = [20, 28, 32, 36, 40, np.inf]
labels = ['20-28', '28-32', '32-36', '36-40', '40+']
data['bmi_group'] = pd.cut(data['bmi'], bins=bins, labels=labels, right=False)

fig, axes = plt.subplots(1, 2, figsize=(16, 7))
# 孕周与Y浓度散点图
scatter1 = axes[0].scatter(data['week'], data['y_concentration'], 
                          c=pd.Categorical(data['bmi_group']).codes, 
                          cmap='viridis', alpha=0.7, s=30)
axes[0].set_xlabel('孕周', fontsize=12)
axes[0].set_ylabel('Y染色体浓度', fontsize=12)
axes[0].set_title('Y染色体浓度 vs 孕周（按BMI分组）', fontsize=14)
cbar = plt.colorbar(scatter1, ax=axes[0])
cbar.set_ticks(range(len(labels)))
cbar.set_ticklabels(labels)
cbar.set_label('BMI分组', fontsize=10)

# BMI与Y浓度散点图
scatter2 = axes[1].scatter(data['bmi'], data['y_concentration'], 
                          c=pd.Categorical(data['bmi_group']).codes, 
                          cmap='viridis', alpha=0.7, s=30)
axes[1].set_xlabel('BMI', fontsize=12)
axes[1].set_ylabel('Y染色体浓度', fontsize=12)
axes[1].set_title('Y染色体浓度 vs BMI（按BMI分组）', fontsize=14)
cbar = plt.colorbar(scatter2, ax=axes[1])
cbar.set_ticks(range(len(labels)))
cbar.set_ticklabels(labels)
cbar.set_label('BMI分组', fontsize=10)

plt.tight_layout()
plt.savefig('scatter_plots.png', dpi=300, bbox_inches='tight')
plt.show()

# 3. 建立模型
print("\n" + "="*60)
print("3. 模型建立与分析")
print("="*60)

# 3.1 简单线性回归模型
print("\n" + "-"*30)
print("3.1 简单线性回归模型")
print("-"*30)

X_linear = data[['week', 'bmi']]
X_linear = sm.add_constant(X_linear)  # 添加常数项
y = data['y_concentration']

model_linear = sm.OLS(y, X_linear).fit()
print(model_linear.summary())

# 3.2 多项式回归模型 (包含二次项和交互项)
print("\n" + "-"*30)
print("3.2 多项式回归模型")
print("-"*30)

# 创建多项式特征
poly = PolynomialFeatures(degree=2, include_bias=False)
X_poly = poly.fit_transform(data[['week', 'bmi']])
feature_names = poly.get_feature_names_out(['week', 'bmi'])

# 转换为DataFrame以便查看系数
X_poly_df = pd.DataFrame(X_poly, columns=feature_names)
X_poly_df = sm.add_constant(X_poly_df)  # 添加常数项

model_poly = sm.OLS(y, X_poly_df).fit()
print(model_poly.summary())

# 3.3 广义可加模型 (GAM)
print("\n" + "-"*30)
print("3.3 广义可加模型 (GAM)")
print("-"*30)

# 使用pygam库
# 定义模型: s()表示平滑项，te()表示张量积平滑(交互作用)
gam = LinearGAM(s(0) + s(1) + te(0, 1)).fit(X_poly[:, :2], y)

# 计算GAM模型的R²
y_pred_gam = gam.predict(X_poly[:, :2])
r2_gam = r2_score(y, y_pred_gam)

# 计算GAM模型的p值
# 使用F检验计算GAM模型的显著性
ssr = np.sum((y_pred_gam - np.mean(y)) ** 2)  # 回归平方和
sse = np.sum((y - y_pred_gam) ** 2)  # 残差平方和
n = len(y)
p = 3  # 模型中的参数数量（三个平滑项）
f_stat = (ssr / p) / (sse / (n - p - 1))
p_value_gam = 1 - stats.f.cdf(f_stat, p, n - p - 1)

# 打印模型摘要
print("GAM模型摘要:")
print(f"R²: {r2_gam:.4f}")
print(f"F统计量: {f_stat:.4f}")
print(f"p值: {p_value_gam:.4e}")

# 4. 模型比较
print("\n" + "="*60)
print("4. 模型比较")
print("="*60)

# 计算各模型的R²
r2_linear = model_linear.rsquared
r2_adj_linear = model_linear.rsquared_adj
r2_poly = model_poly.rsquared
r2_adj_poly = model_poly.rsquared_adj

print(f"线性模型 R²: {r2_linear:.4f}, 调整R²: {r2_adj_linear:.4f}, p值: {model_linear.f_pvalue:.4e}")
print(f"多项式模型 R²: {r2_poly:.4f}, 调整R²: {r2_adj_poly:.4f}, p值: {model_poly.f_pvalue:.4e}")
print(f"GAM模型 R²: {r2_gam:.4f}, p值: {p_value_gam:.4e}")

# 5. 可视化模型拟合效果
print("\n" + "="*60)
print("5. 模型可视化")
print("="*60)

fig, axes = plt.subplots(2, 2, figsize=(16, 14))

# 5.1 实际值与预测值散点图 - 线性模型
y_pred_linear = model_linear.predict(X_linear)
axes[0, 0].scatter(y, y_pred_linear, alpha=0.6, s=30)
axes[0, 0].plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=2)
axes[0, 0].set_xlabel('实际值', fontsize=12)
axes[0, 0].set_ylabel('预测值', fontsize=12)
axes[0, 0].set_title('线性模型: 实际值 vs 预测值', fontsize=14)
axes[0, 0].text(0.05, 0.95, f'$R^2$ = {r2_linear:.4f}\np = {model_linear.f_pvalue:.2e}', 
                transform=axes[0, 0].transAxes, fontsize=12,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

# 5.2 实际值与预测值散点图 - 多项式模型
y_pred_poly = model_poly.predict(X_poly_df)
axes[0, 1].scatter(y, y_pred_poly, alpha=0.6, s=30)
axes[0, 1].plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=2)
axes[0, 1].set_xlabel('实际值', fontsize=12)
axes[0, 1].set_ylabel('预测值', fontsize=12)
axes[0, 1].set_title('多项式模型: 实际值 vs 预测值', fontsize=14)
axes[0, 1].text(0.05, 0.95, f'$R^2$ = {r2_poly:.4f}\np = {model_poly.f_pvalue:.2e}', 
                transform=axes[0, 1].transAxes, fontsize=12,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

# 5.3 实际值与预测值散点图 - GAM模型
axes[1, 0].scatter(y, y_pred_gam, alpha=0.6, s=30)
axes[1, 0].plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=2)
axes[1, 0].set_xlabel('实际值', fontsize=12)
axes[1, 0].set_ylabel('预测值', fontsize=12)
axes[1, 0].set_title('GAM模型: 实际值 vs 预测值', fontsize=14)
axes[1, 0].text(0.05, 0.95, f'$R^2$ = {r2_gam:.4f}\np = {p_value_gam:.2e}', 
                transform=axes[1, 0].transAxes, fontsize=12,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

# 5.4 残差图 - 多项式模型
residuals = y - y_pred_poly
axes[1, 1].scatter(y_pred_poly, residuals, alpha=0.6, s=30)
axes[1, 1].axhline(y=0, color='r', linestyle='--', lw=2)
axes[1, 1].set_xlabel('预测值', fontsize=12)
axes[1, 1].set_ylabel('残差', fontsize=12)
axes[1, 1].set_title('多项式模型: 残差图', fontsize=14)

plt.tight_layout()
plt.savefig('model_comparison.png', dpi=300, bbox_inches='tight')
plt.show()

# 6. GAM模型的平滑效应可视化
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# 孕周的平滑效应
XX = gam.generate_X_grid(term=0)
axes[0].plot(XX[:, 0], gam.partial_dependence(term=0, X=XX), 'b-', lw=2)
axes[0].scatter(data['week'], [0] * len(data), alpha=0.1)
axes[0].set_title('孕周对Y染色体浓度的平滑效应', fontsize=14)
axes[0].set_xlabel('孕周', fontsize=12)
axes[0].set_ylabel('偏依赖', fontsize=12)
axes[0].grid(True, alpha=0.3)

# BMI的平滑效应
XX = gam.generate_X_grid(term=1)
axes[1].plot(XX[:, 1], gam.partial_dependence(term=1, X=XX), 'g-', lw=2)
axes[1].scatter(data['bmi'], [0] * len(data), alpha=0.1)
axes[1].set_title('BMI对Y染色体浓度的平滑效应', fontsize=14)
axes[1].set_xlabel('BMI', fontsize=12)
axes[1].set_ylabel('偏依赖', fontsize=12)
axes[1].grid(True, alpha=0.3)

# 交互效应的等高线图
# 生成网格数据
n_points = 30
xx, yy = np.meshgrid(np.linspace(data['week'].min(), data['week'].max(), n_points),
                     np.linspace(data['bmi'].min(), data['bmi'].max(), n_points))
ZZ = np.zeros(xx.shape)
for i in range(n_points):
    for j in range(n_points):
        ZZ[j, i] = gam.partial_dependence(term=2, X=np.array([[xx[i, j], yy[i, j]]]))
        
contour = axes[2].contour(xx, yy, ZZ, levels=14, cmap='viridis')
axes[2].clabel(contour, inline=True, fontsize=8)
axes[2].set_title('孕周和BMI的交互效应', fontsize=14)
axes[2].set_xlabel('孕周', fontsize=12)
axes[2].set_ylabel('BMI', fontsize=12)

plt.tight_layout()
plt.savefig('gam_smooth_effects.png', dpi=300, bbox_inches='tight')
plt.show()

# 7. 多重共线性诊断与处理
print("\n" + "="*60)
print("7. 多重共线性诊断与处理")
print("="*60)

# 计算方差膨胀因子(VIF)来诊断多重共线性
vif_data = pd.DataFrame()
vif_data["feature"] = X_poly_df.columns
vif_data["VIF"] = [variance_inflation_factor(X_poly_df.values, i) 
                   for i in range(len(X_poly_df.columns))]

print("多项式特征的方差膨胀因子(VIF):")
print(vif_data)

# 对变量进行中心化处理以减少多重共线性
data_centered = data.copy()
data_centered['week_centered'] = data['week'] - data['week'].mean()
data_centered['bmi_centered'] = data['bmi'] - data['bmi'].mean()

# 使用中心化后的变量重新运行多项式回归
poly_centered = PolynomialFeatures(degree=2, include_bias=False)
X_poly_centered = poly_centered.fit_transform(data_centered[['week_centered', 'bmi_centered']])
feature_names_centered = poly.get_feature_names_out(['week', 'bmi'])

X_poly_df_centered = pd.DataFrame(X_poly_centered, columns=feature_names_centered)
X_poly_df_centered = sm.add_constant(X_poly_df_centered)

model_poly_centered = sm.OLS(y, X_poly_df_centered).fit()
print("\n中心化后的多项式模型系数:")
print(model_poly_centered.params)
print("\n中心化后的多项式模型p值:")
print(model_poly_centered.pvalues)

# 计算孕周的边际效应
# 孕周的边际效应 = β1 + 2*β3*week + β5*bmi
week_mean = data['week'].mean()
bmi_mean = data['bmi'].mean()

# 使用原始模型系数
marginal_effect_week = (model_poly.params['week'] + 
                        2 * model_poly.params['week^2'] * week_mean + 
                        model_poly.params['week bmi'] * bmi_mean)

print(f"\n孕周在平均值处的边际效应: {marginal_effect_week:.6f}")

# 绘制孕周与Y染色体浓度的偏回归图
plt.figure(figsize=(10, 6))
# 控制BMI不变，绘制孕周与Y染色体浓度的关系
bmi_groups = data['bmi_group'].unique()
colors = plt.cm.viridis(np.linspace(0, 1, len(bmi_groups)))

for i, group in enumerate(bmi_groups):
    group_data = data[data['bmi_group'] == group]
    if len(group_data) > 5:  # 确保每组有足够样本
        # 计算该BMI组内孕周与Y浓度的关系
        X_temp = group_data[['week']]
        X_temp = sm.add_constant(X_temp)
        model_temp = sm.OLS(group_data['y_concentration'], X_temp).fit()
        
        # 绘制散点图和拟合线
        plt.scatter(group_data['week'], group_data['y_concentration'], 
                   color=colors[i], alpha=0.6, label=f'BMI {group}', s=30)
        
        # 绘制拟合线
        x_range = np.linspace(group_data['week'].min(), group_data['week'].max(), 100)
        X_pred = sm.add_constant(x_range)
        y_pred = model_temp.predict(X_pred)
        plt.plot(x_range, y_pred, color=colors[i], linestyle='-', linewidth=2)

plt.xlabel('孕周')
plt.ylabel('Y染色体浓度')
plt.title('不同BMI分组下孕周与Y染色体浓度的关系')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('week_vs_yconcentration_by_bmi.png', dpi=300, bbox_inches='tight')
plt.show()

# 8. 输出模型系数和显著性
print("\n" + "="*60)
print("8. 模型系数和显著性")
print("="*60)

print("线性模型系数:")
print(model_linear.params)
print("\n线性模型p值:")
print(model_linear.pvalues)

print("\n多项式模型系数:")
print(model_poly.params)
print("\n多项式模型p值:")
print(model_poly.pvalues)

print("\n中心化后的多项式模型系数:")
print(model_poly_centered.params)
print("\n中心化后的多项式模型p值:")
print(model_poly_centered.pvalues)

# 9. 保存结果到文件
print("\n" + "="*60)
print("9. 结果保存")
print("="*60)

# 保存模型比较结果
results_df = pd.DataFrame({
    '模型': ['线性回归', '多项式回归', 'GAM', '中心化多项式回归'],
    'R²': [r2_linear, r2_poly, r2_gam, model_poly_centered.rsquared],
    '调整R²': [r2_adj_linear, r2_adj_poly, np.nan, model_poly_centered.rsquared_adj],
    'p值': [model_linear.f_pvalue, model_poly.f_pvalue, p_value_gam, model_poly_centered.f_pvalue]
})

results_df.to_csv('model_results.csv', index=False, encoding='utf-8-sig')
print("模型比较结果已保存到 'model_results.csv'")

# 保存多项式模型系数
coef_df = pd.DataFrame({
    '变量': model_poly.params.index,
    '系数': model_poly.params.values,
    '标准误': model_poly.bse.values,
    't值': model_poly.tvalues.values,
    'p值': model_poly.pvalues.values
})
coef_df.to_csv('polynomial_coefficients.csv', index=False, encoding='utf-8-sig')
print("多项式模型系数已保存到 'polynomial_coefficients.csv'")

# 保存中心化多项式模型系数
coef_centered_df = pd.DataFrame({
    '变量': model_poly_centered.params.index,
    '系数': model_poly_centered.params.values,
    '标准误': model_poly_centered.bse.values,
    't值': model_poly_centered.tvalues.values,
    'p值': model_poly_centered.pvalues.values
})
coef_centered_df.to_csv('polynomial_coefficients_centered.csv', index=False, encoding='utf-8-sig')
print("中心化多项式模型系数已保存到 'polynomial_coefficients_centered.csv'")

# 保存相关系数矩阵
corr_matrix.to_csv('correlation_matrix.csv', encoding='utf-8-sig')
print("相关系数矩阵已保存到 'correlation_matrix.csv'")

# 保存VIF结果
vif_data.to_csv('vif_results.csv', index=False, encoding='utf-8-sig')
print("VIF结果已保存到 'vif_results.csv'")

print("\n分析完成！所有图表和结果已保存。")

# 10. 生成报告摘要
print("\n" + "="*60)
print("10. 分析报告摘要")
print("="*60)

print(f"数据集大小: {data.shape[0]}个样本")
print(f"Y染色体浓度与孕周的相关系数: {corr_matrix.loc['week', 'y_concentration']:.3f}")
print(f"Y染色体浓度与BMI的相关系数: {corr_matrix.loc['bmi', 'y_concentration']:.3f}")
print(f"最佳模型: 多项式回归模型 (R² = {r2_poly:.4f}, p = {model_poly.f_pvalue:.2e})")
print(f"孕周对Y染色体浓度的影响: {'正向' if model_poly.params['week'] > 0 else '负向'} (p = {model_poly.pvalues['week']:.3f})")
print(f"BMI对Y染色体浓度的影响: {'正向' if model_poly.params['bmi'] > 0 else '负向'} (p = {model_poly.pvalues['bmi']:.3f})")
print(f"孕周二次项显著性: {'显著' if model_poly.pvalues['week^2'] < 0.05 else '不显著'} (p = {model_poly.pvalues['week^2']:.3f})")
print(f"BMI二次项显著性: {'显著' if model_poly.pvalues['bmi^2'] < 0.05 else '不显著'} (p = {model_poly.pvalues['bmi^2']:.3f})")
print(f"交互项显著性: {'显著' if model_poly.pvalues['week bmi'] < 0.05 else '不显著'} (p = {model_poly.pvalues['week bmi']:.3f})")
print(f"孕周在平均值处的边际效应: {marginal_effect_week:.6f}")
